from django.db import models
from dvadmin.utils.models import CoreModel


class Projects(CoreModel):
    TYPE_CHOICES = (
        (0, "图像分类"),
        (1, "目标检测"),
        (2, "图像分割"),
    )
    type = models.IntegerField(
        choices=TYPE_CHOICES, default=0, verbose_name="项目类型", null=True, blank=True, help_text="项目类型"
    )
    name = models.CharField(max_length=255, verbose_name="项目名称", null=True, blank=True, help_text="项目名称")
    description = models.CharField(max_length=255, verbose_name="项目描述", null=True, blank=True, help_text="项目描述")

    def __str__(self):
        return self.name

    class Meta:
        db_table = "project"
        verbose_name = '项目信息表'
        verbose_name_plural = verbose_name
        ordering = ('-create_datetime',)


class Models(CoreModel):
    TYPE_CHOICES = (
        (0, "图像分类"),
        (1, "目标检测"),
        (2, "图像分割"),
    )
    type = models.IntegerField(
        choices=TYPE_CHOICES, default=0, verbose_name="模型类型", null=True, blank=True, help_text="模型类型"
    )
    name = models.CharField(max_length=255, verbose_name="模型名称", null=True, blank=True, help_text="模型名称")
    version = models.CharField(max_length=10, verbose_name="模型版本", null=True, blank=True, help_text="模型版本")
    project = models.ForeignKey(to="Projects", null=True, blank=True, verbose_name='所属项目ID', help_text="所属项目ID",
                                on_delete=models.CASCADE, db_constraint=False, related_name='children')
    dataset_id = models.IntegerField(verbose_name="数据集ID", null=True, blank=True, help_text="数据集ID")
    STATUS_CHOICES = (
        (0, "未训练"),
        (1, "正在训练"),
        (2, "训练完成"),
        (3, "训练终止"),
    )
    status = models.IntegerField(
        choices=STATUS_CHOICES, default=0, verbose_name="训练状态", null=True, blank=True, help_text="训练状态"
    )

    time = models.CharField(max_length=255, verbose_name="训练时长", null=True, blank=True, help_text="训练时长")
    weights = models.CharField(max_length=255, verbose_name="模型权重路径", null=True, blank=True, help_text="模型权重路径")
    cfg = models.CharField(max_length=255, verbose_name="配置文件路径", null=True, blank=True, help_text="配置文件路径")
    epochs = models.IntegerField(default=100, verbose_name="循环次数", null=True, blank=True, help_text="循环次数")
    result = models.CharField(max_length=255, verbose_name="结果路径", null=True, blank=True, help_text="结果路径")
    model_version = models.CharField(max_length=255, verbose_name="模型版本", null=True, blank=True, help_text="模型版本")
    is_published = models.BooleanField(verbose_name="是否已发布", help_text='是否已发布', default=False, db_index=True)

    def __str__(self):
        return self.name

    class Meta:
        db_table = "model"
        verbose_name = '模型信息表'
        verbose_name_plural = verbose_name
        ordering = ('-create_datetime',)


class Services(CoreModel):
    TYPE_CHOICES = (
        (0, "图像分类"),
        (1, "目标检测"),
        (2, "图像分割"),
    )
    type = models.IntegerField(
        choices=TYPE_CHOICES, default=0, verbose_name="服务类型", null=True, blank=True, help_text="服务类型"
    )
    name = models.CharField(max_length=255, verbose_name="服务名称", null=True, blank=True, help_text="服务名称")
    project = models.ForeignKey(to="Projects", null=True, blank=True, verbose_name='所属项目ID', help_text="所属项目ID",
                                on_delete=models.CASCADE, db_constraint=False)
    model = models.ForeignKey(to="Models", null=True, blank=True, verbose_name='所属模型ID', help_text="所属模型ID",
                              on_delete=models.CASCADE, db_constraint=False)
    STATUS_CHOICES = (
        (0, "下线"),
        (1, "上线"),
    )
    status = models.IntegerField(
        choices=STATUS_CHOICES, default=0, verbose_name="服务状态", null=True, blank=True, help_text="服务状态"
    )

    class Meta:
        db_table = "service"
        verbose_name = '服务信息表'
        verbose_name_plural = verbose_name
        ordering = ('-create_datetime',)
